import os
from typing import List, Optional, Tuple, Type, Union

import aiohttp
import requests

from lagent.schema import ActionReturn, ActionStatusCode
from .base_action import AsyncActionMixin, BaseAction, tool_api
from .parser import BaseParser, JsonParser


class GoogleSearch(BaseAction):
    """
    Serper.dev 谷歌搜索 API 的封装类。
    基于 langchain 的 GoogleSerperAPIWrapper 修改。
    
    Args:
        api_key: Serper API 密钥，可在 https://serper.dev 获取
        timeout: Serper 请求的超时时间
        search_type: 搜索类型，支持 ['search', 'images', 'news', 'places']
        description: 工具描述
        parser: 结果解析器类
    """

    # 不同搜索类型对应的结果键名
    result_key_for_type = {
        'news': 'news',
        'places': 'places', 
        'images': 'images',
        'search': 'organic',
    }

    def __init__(
        self,
        api_key: Optional[str] = None,
        timeout: int = 5,
        search_type: str = 'search',
        description: Optional[dict] = None,
        parser: Type[BaseParser] = JsonParser,
    ):
        super().__init__(description, parser)
        # 从环境变量或参数获取 API 密钥
        api_key = os.environ.get('SERPER_API_KEY', api_key)
        if api_key is None:
            raise ValueError('请设置 Serper API 密钥(环境变量 SERPER_API_KEY 或参数 api_key)')
        self.api_key = api_key
        self.timeout = timeout
        self.search_type = search_type

    @tool_api
    def run(self, query: str, k: int = 10) -> ActionReturn:
        """
        执行谷歌搜索。适用于需要简短明了答案的问题。

        Args:
            query: 搜索查询
            k: 返回前k个搜索结果
        """
        tool_return = ActionReturn(type=self.name)
        # 执行搜索
        status_code, response = self._search(query, k=k)
        
        # 处理不同状态的返回
        if status_code == -1:  # 请求错误
            tool_return.errmsg = response
            tool_return.state = ActionStatusCode.HTTP_ERROR
        elif status_code == 200:  # 请求成功
            parsed_res = self._parse_results(response, k)
            tool_return.result = [dict(type='text', content=str(parsed_res))]
            tool_return.state = ActionStatusCode.SUCCESS
        else:  # API 错误
            tool_return.errmsg = str(status_code)
            tool_return.state = ActionStatusCode.API_ERROR
        return tool_return

    def _parse_results(self, results: dict, k: int) -> Union[str, List[str]]:
        """
        解析 Serper API 的搜索结果。

        Args:
            results: Serper API 返回的 JSON 格式结果
            k: 返回结果数量

        Returns:
            解析后的搜索结果列表
        """
        snippets = []

        # 处理直接答案框
        if results.get('answerBox'):
            answer_box = results.get('answerBox', {})
            if answer_box.get('answer'):
                return [answer_box.get('answer')]
            elif answer_box.get('snippet'):
                return [answer_box.get('snippet').replace('\n', ' ')]
            elif answer_box.get('snippetHighlighted'):
                return answer_box.get('snippetHighlighted')

        # 处理知识图谱
        if results.get('knowledgeGraph'):
            kg = results.get('knowledgeGraph', {})
            title = kg.get('title')
            entity_type = kg.get('type')
            if entity_type:
                snippets.append(f'{title}: {entity_type}.')
            description = kg.get('description')
            if description:
                snippets.append(description)
            for attribute, value in kg.get('attributes', {}).items():
                snippets.append(f'{title} {attribute}: {value}.')

        # 处理搜索结果
        for result in results[self.result_key_for_type[self.search_type]][:k]:
            if 'snippet' in result:
                snippets.append(result['snippet'])
            for attribute, value in result.get('attributes', {}).items():
                snippets.append(f'{attribute}: {value}.')

        return snippets if snippets else ['No good Google Search Result was found']

    def _search(self, search_term: str, search_type: Optional[str] = None, **kwargs) -> Tuple[int, Union[dict, str]]:
        """
        向 Serper API 发送 HTTP 请求。

        Args:
            search_term: 搜索查询
            search_type: 搜索类型，默认为 'search'

        Returns:
            tuple: (状态码, 响应内容)
        """
        # 构建请求头和参数
        headers = {
            'X-API-KEY': self.api_key or '',
            'Content-Type': 'application/json',
        }
        params = {
            'q': search_term,
            **{key: value for key, value in kwargs.items() if value is not None},
        }
        
        try:
            # 发送 POST 请求
            response = requests.post(
                f'https://google.serper.dev/{search_type or self.search_type}',
                headers=headers,
                params=params,
                timeout=self.timeout)
        except Exception as e:
            return -1, str(e)
        return response.status_code, response.json()


class AsyncGoogleSearch(AsyncActionMixin, GoogleSearch):
    """异步版本的谷歌搜索类"""
    
    # 异步版本的方法实现与同步版本类似，
    # 主要区别是使用 aiohttp 进行异步 HTTP 请求
    # 这里省略了重复的注释...
